from tqdm import tqdm

from zkl_knowedit_ft import FinetuneCallback, FinetuneHparams, FinetuneMetrics


class TqdmFinetuneCallback(FinetuneCallback):
    def __init__(self, desc: str | None = "Finetuning"):
        self._desc: str | None = desc
        self._tqdm: tqdm | None = None

    def on_start(self, hparams: FinetuneHparams):
        self._tqdm = tqdm(unit='epoch', total=hparams.stopping_epochs_num, desc=self._desc)

    def on_epoch_start(self, epoch_i: int):
        pass

    def on_epoch_stop(self, metrics: FinetuneMetrics):
        if isinstance(self._tqdm, tqdm):
            self._tqdm.update(1)
            self._tqdm.set_postfix({
                "rewriting_ce": metrics.rewriting_ce,
                "rewriting_acc": metrics.rewriting_acc})

    def on_stop(self, metrics: FinetuneMetrics | None):
        if self._tqdm is not None:
            self._tqdm.close()
